package middleware

import (
	"github.com/gin-gonic/gin"
	"go.uber.org/ratelimit"
	"golang.org/x/time/rate"
	"net/http"
)

// GetLeakyRateLimitMiddleware 漏斗算法
func GetLeakyRateLimitMiddleware(qps int) func(ctx *gin.Context) {
	funnel := ratelimit.New(qps)
	return func(ctx *gin.Context) {
		_ = funnel.Take()
	}
}

// GetTokenRateLimitMiddleware 令牌桶算法
func GetTokenRateLimitMiddleware(qps int, errorResponse interface{}) func(ctx *gin.Context) {
	bucket := rate.NewLimiter(rate.Limit(qps), qps)
	return func(ctx *gin.Context) {
		if !bucket.Allow() {
			ctx.JSON(http.StatusOK, errorResponse)
			ctx.Abort()
		}
	}
}
